package com.jd.platform.async.executor;

import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

/**
 * TraceId 上下文管理工具
 * 
 * 用于在多线程异步任务中传递 TraceId，解决日志追踪问题
 * 
 * @author AsyncTool
 * @version 1.5.0
 */
public class TraceContext {
    
    /**
     * MDC 中 TraceId 的 key
     */
    public static final String TRACE_ID_KEY = "traceId";
    
    /**
     * MDC 中 WorkflowId 的 key（可选，用于区分不同的任务链）
     */
    public static final String WORKFLOW_ID_KEY = "workflowId";
    
    /**
     * MDC 中 WorkerId 的 key（可选，用于区分具体的任务）
     */
    public static final String WORKER_ID_KEY = "workerId";
    
    /**
     * 生成一个新的 TraceId
     * 格式：trace_timestamp_random
     * 
     * @return 新的 TraceId
     */
    public static String generateTraceId() {
        return "trace_" + System.currentTimeMillis() + "_" + 
               UUID.randomUUID().toString().substring(0, 8);
    }
    
    /**
     * 设置当前线程的 TraceId
     * 
     * @param traceId TraceId
     */
    public static void setTraceId(String traceId) {
        if (traceId != null && !traceId.isEmpty()) {
            MDC.put(TRACE_ID_KEY, traceId);
        }
    }
    
    /**
     * 获取当前线程的 TraceId
     * 
     * @return TraceId，如果没有则返回 null
     */
    public static String getTraceId() {
        return MDC.get(TRACE_ID_KEY);
    }
    
    /**
     * 设置 WorkflowId（任务链 ID）
     * 
     * @param workflowId 任务链 ID
     */
    public static void setWorkflowId(String workflowId) {
        if (workflowId != null && !workflowId.isEmpty()) {
            MDC.put(WORKFLOW_ID_KEY, workflowId);
        }
    }
    
    /**
     * 获取 WorkflowId
     * 
     * @return WorkflowId
     */
    public static String getWorkflowId() {
        return MDC.get(WORKFLOW_ID_KEY);
    }
    
    /**
     * 设置 WorkerId（具体任务 ID）
     * 
     * @param workerId 任务 ID
     */
    public static void setWorkerId(String workerId) {
        if (workerId != null && !workerId.isEmpty()) {
            MDC.put(WORKER_ID_KEY, workerId);
        }
    }
    
    /**
     * 获取 WorkerId
     * 
     * @return WorkerId
     */
    public static String getWorkerId() {
        return MDC.get(WORKER_ID_KEY);
    }
    
    /**
     * 清除当前线程的所有 TraceContext
     */
    public static void clear() {
        MDC.remove(TRACE_ID_KEY);
        MDC.remove(WORKFLOW_ID_KEY);
        MDC.remove(WORKER_ID_KEY);
    }
    
    /**
     * 清除 TraceId
     */
    public static void clearTraceId() {
        MDC.remove(TRACE_ID_KEY);
    }
    
    /**
     * 获取当前线程的所有 MDC 上下文（用于传递到子线程）
     * 
     * @return MDC 上下文的副本
     */
    public static Map<String, String> getMDCContext() {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return contextMap != null ? contextMap : new HashMap<>();
    }
    
    /**
     * 设置 MDC 上下文（用于从父线程传递过来）
     * 
     * @param contextMap MDC 上下文
     */
    public static void setMDCContext(Map<String, String> contextMap) {
        if (contextMap != null && !contextMap.isEmpty()) {
            MDC.setContextMap(contextMap);
        }
    }
    
    /**
     * 如果当前线程没有 TraceId，则自动生成一个
     * 
     * @return 当前 TraceId（可能是新生成的）
     */
    public static String ensureTraceId() {
        String traceId = getTraceId();
        if (traceId == null || traceId.isEmpty()) {
            traceId = generateTraceId();
            setTraceId(traceId);
        }
        return traceId;
    }
}

